import cv2
import numpy as np

def mask_to_points(label_mask, alpha = 0.01):
        
    contours, hierarchy = cv2.findContours(label_mask,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_SIMPLE)
    w_h = label_mask.shape[::-1]
    points=[]
    # temp_mask = np.zeros_like(label_mask)
    for idx, contour in enumerate(contours): 
        # 轮廓近似
        if alpha > 0:
            epsilon = alpha*cv2.arcLength(contour,True)
            contour = cv2.approxPolyDP(contour,epsilon,True)
        contour = np.squeeze(contour,1)/w_h
        points.append(contour.tolist())
    return points

class MaskToPoints():
    def __init__(self, keys, cls_num=2, alpha = 0.01):
        self.cls_num = cls_num
        self.keys = keys
        self.alpha = alpha

    def __call__(self, data:dict):
        seg_result = data[self.keys["in"]]
        total_points = []
        for idx_, mask in enumerate(seg_result.get("masks",[])):
            cur_img_points = []
            for cls_id in range(self.cls_num): # 逐个类处理掩码结果
                cur_mask = np.zeros_like(mask,np.uint8)
                cur_mask[mask == cls_id] = 255
                cur_mask = cv2.morphologyEx(cur_mask,cv2.MORPH_OPEN, np.ones((5,5),np.uint8)) # 开运算
                # 掩码图转点集
                points = mask_to_points(cur_mask,self.alpha)
                cur_img_points.append(points)
            total_points.append(cur_img_points)
        if self.keys.get("out",False):
            data[self.keys["out"]]["points"] = total_points
        else:
            data[self.keys["in"]]["points"] = total_points
        return data